#include <cstring>
#include <iostream>
#include <faiss/ascend/AscendIndexFlat.h>
#include "ascend_index.h"

using namespace std;

namespace {
    faiss::ascend::AscendIndexFlat *_index = nullptr;
    int _dim = 0;
    int _device_id = 0;
}

void init(int device_id, int dim) {
    faiss::ascend::AscendIndexFlatConfig conf({device_id});
    _index = new faiss::ascend::AscendIndexFlat(dim, faiss::METRIC_L2, conf);
    _dim = dim;
    _device_id = device_id;
}

void insert_gallery(int baseSize, float *feat) {
    _index->add(baseSize, feat);
}

void search(int n, float *queryFeat, int topK, float *scores, int64_t *idxs) {
    _index->search((faiss::Index::idx_t) n, queryFeat, (faiss::Index::idx_t) topK, scores, idxs);
}

int getBaseSize() {
    return _index->ntotal;
}